import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import akshare as ak
import time


# 下载数据，添加重试逻辑
def download_data(ticker, start_date, end_date, max_retries=3):
    retries = 0
    while retries < max_retries:
        try:
            # 使用akshare获取股票日线数据
            data = ak.stock_zh_a_hist(symbol=ticker, period="daily", start_date=start_date, end_date=end_date)
            if not data.empty:
                # 处理数据，将日期列转换为日期类型并设置为索引
                data['日期'] = pd.to_datetime(data['日期'])
                data.set_index('日期', inplace=True)
                # 重命名列名以符合原代码逻辑
                data.rename(columns={'开盘': 'Open', '最高': 'High', '最低': 'Low', '收盘': 'Close', '成交量': 'Volume'}, inplace=True)
                return data
        except Exception as e:
            print(f"第 {retries + 1} 次下载出错: {e}")
        retries += 1
        time.sleep(2)  # 等待 2 秒后重试
    print("多次尝试下载均失败，请检查网络或数据源。")
    return None


# 交易信号生成
def generate_signals(data, short_window=5, long_window=20):
    data['short_mavg'] = data['Close'].rolling(window=short_window).mean()
    data['long_mavg'] = data['Close'].rolling(window=long_window).mean()
    data['signal'] = 0
    data.loc[data['short_mavg'] > data['long_mavg'], 'signal'] = 1
    data.loc[data['short_mavg'] < data['long_mavg'], 'signal'] = -1
    data['position'] = data['signal'].diff()
    return data


# 止损与止盈
def stop_loss_profit(data, stop_loss=0.05, take_profit=0.1):
    positions = []
    current_position = 0
    entry_price = 0
    for i in range(len(data)):
        if data['position'].iloc[i] == 1:
            current_position = 1
            entry_price = data['Close'].iloc[i]
        elif data['position'].iloc[i] == -1:
            current_position = 0
        elif current_position == 1:
            current_price = data['Close'].iloc[i]
            if (current_price <= entry_price * (1 - stop_loss)) or (current_price >= entry_price * (1 + take_profit)):
                current_position = 0
        positions.append(current_position)
    data['position_after_stop'] = positions
    return data


# 回测
def backtest(data):
    data['returns'] = data['Close'].pct_change()
    data['strategy_returns'] = data['position_after_stop'].shift(1) * data['returns']
    cumulative_returns = (1 + data['strategy_returns']).cumprod()
    return cumulative_returns


# 可视化
def visualize(data, cumulative_returns):
    plt.figure(figsize=(12, 8))
    plt.subplot(2, 1, 1)
    plt.plot(data['Close'], label='Close Price')
    plt.plot(data['short_mavg'], label='Short MA')
    plt.plot(data['long_mavg'], label='Long MA')
    plt.plot(data[data['position'] == 1].index, data['short_mavg'][data['position'] == 1], '^', markersize=10, color='g',
             label='Buy Signal')
    plt.plot(data[data['position'] == -1].index, data['short_mavg'][data['position'] == -1], 'v', markersize=10, color='r',
             label='Sell Signal')
    plt.title('Stock Price and Moving Averages')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()

    plt.subplot(2, 1, 2)
    plt.plot(cumulative_returns, label='Strategy Cumulative Returns')
    plt.title('Strategy Cumulative Returns')
    plt.xlabel('Date')
    plt.ylabel('Cumulative Returns')
    plt.legend()

    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    # 选择百利天恒，akshare中股票代码不需要后缀
    ticker = "688506"
    start_date = '20230101'
    end_date = '20231231'

    # 下载数据
    data = download_data(ticker, start_date, end_date)
    if data is not None:
        # 生成交易信号
        data = generate_signals(data)
        # 止损与止盈
        data = stop_loss_profit(data)
        # 回测
        cumulative_returns = backtest(data)
        # 可视化
        visualize(data, cumulative_returns)